import sys
import pickle
import pandas as pd

#Import and unpickle relevant data
files = sys.argv[1:]
unpickle =[[] for file in range(len(files))]

for file_num, file in enumerate(files):
	f = open(file,'rb')
	unpickle[file_num] = pickle.load(f)
	f.close()
	
time = unpickle[0]
event_list = unpickle[1]
onset_dict = unpickle[2]
event_num_check = unpickle[3]

#Trim events to window signal from -1s to 4s and -1s to 5s
key_names = ['5s Trimmed Events','6s Trimmed Events']
trimmed_events = {name:[] for name in key_names}
temp_events = [[] for number in range(2)]

for event_num, event in enumerate(event_list):
	for item_num, item in enumerate(event):
		if time[item_num] >= -1 and time[item_num] <= 4:
			temp_events[0].append(item)
		if time[item_num] >= -1 and time[item_num] <= 5:
			temp_events[1].append(item)
	trimmed_events['5s Trimmed Events'].append(temp_events[0])
	trimmed_events['6s Trimmed Events'].append(temp_events[1])
	temp_events = [[] for number in range(2)]
	

#Calculate individual event means and overall event mean and SEM for 5 and 6s signal
calc_dict = {'5s Mean':[],'5s SEM':[],'6s Mean':[],'6s SEM':[]}

for name in key_names:
	event_calc = pd.DataFrame(trimmed_events[name])
	event_calc = event_calc.T
	event_calc_all_avg = event_calc.mean(axis = 0, numeric_only = True)
	event_calc_sem = event_calc_all_avg.sem(axis = 0)
	event_calc_avg = event_calc_all_avg.mean(axis = 0)
	if name == '5s Trimmed Events':
		onset_dict['5s Average'] = event_calc_all_avg.tolist()
		if event_num_check > 1:
			calc_dict['5s Mean'] = event_calc_avg
			calc_dict['5s SEM'] = event_calc_sem
		else:
			calc_dict['5s Mean'] = 'nan'
			calc_dict['5s SEM'] = 'nan'
	if name == '6s Trimmed Events':
		onset_dict['6s Average'] = event_calc_all_avg.tolist()
		if event_num_check > 1:
			calc_dict['6s Mean'] = event_calc_avg
			calc_dict['6s SEM'] = event_calc_sem
		else:
			calc_dict['6s Mean'] = 'nan'
			calc_dict['6s SEM'] = 'nan'
		
#Save and pickle calculations	
pickle.dump(calc_dict, open('signal_mean_dict.pkl','wb'))
pickle.dump(onset_dict, open('onset_dict.pkl','wb'))